# -*- coding: utf-8 -*-

import numpy as np
import logging


def load_double_moon(length=200, dst=-3, rad=10, width=2):
    """generate double moon data

    Parameters
    ----------
    length : 待生成的数据的长度
    dst : 圆心距
    rad : 半径
    width : 数据带宽度
    """
    N1 = 10 * length
    w2 = width / 2
    done = True
    data = np.empty(0)
    while done:
        # generate Rectangular data
        tmp_x = 2 * (rad + w2) * (np.random.random([N1, 1]) - 0.5)
        tmp_y = (rad + w2) * np.random.random([N1, 1])
        tmp = np.concatenate((tmp_x, tmp_y), axis=1)
        tmp_ds = np.sqrt(tmp_x * tmp_x + tmp_y * tmp_y)
        # generate double moon data ---upper
        idx = np.logical_and(tmp_ds > (rad - w2), tmp_ds < (rad + w2))
        idx = (idx.nonzero())[0]

        if data.shape[0] == 0:
            data = tmp.take(idx, axis=0)
        else:
            data = np.concatenate((data, tmp.take(idx, axis=0)), axis=0)
        if data.shape[0] >= length:
            done = False
    db_moon = data[0:length, :]
    if __name__ == "__main__":
        print(data)
        print(db_moon)
    # generate double moon data ----down
    data_t = np.empty([length, 2])
    data_t[:, 0] = data[0:length, 0] + rad
    data_t[:, 1] = -data[0:length, 1] - dst
    db_moon = np.concatenate((db_moon, data_t), axis=0)
    logging.warning("db_moon shape: " + str(db_moon.shape))
    return db_moon


if __name__ == "__main__":
    import matplotlib.pyplot as plt

    n = 20
    d = -3
    r = 10
    w = 2

    data = load_double_moon(n, d, r, w)

    plt.plot(data[0:n, 0], data[0:n, 1], 'r+', data[n:2 * n, 0], data[n:2 * n, 1], 'r+')
    plt.show()